# -*- coding: utf-8 -*-
"""
Created on Sat Jun 12 14:39:29 2021

@author: zhuo 木鸟
"""

import pandas as pd
import numpy as np
from autoregression_analyze import evalute_autoregression
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
from series_to_supervised import series_to_supervised
from analyze_series import acf_plotting, pacf_plotting

plt.rcParams['font.sans-serif']=['SimHei']    #画图时使用中文字体
plt.rcParams['axes.unicode_minus'] = False

def plot_residual_error(time_range, errors):
    # 画出残差图
    fig = plt.figure()
    plt.plot(time_range, errors)
    plt.xlabel('时间', fontsize=16)
    plt.ylabel('残差', fontsize=16)
    plt.title('残差图', fontsize=16)
    plt.grid()
    fig.savefig('../图片/残差原始数据.png')
    plt.show()
    

if __name__ == '__main__':
    path = r'../附件/中间数据/光伏-建筑板块市值.xlsx'
    series = pd.read_excel(path, header=0, index_col=0, parse_dates=True, squeeze=True)
    series_window = series_to_supervised(series, n_in=2, n_out=1, dropnan=True)
    
    # 拆分数据集
    train, test = series_window['2019-4-1':'2021-4-30'], series_window['2021-5-1':'2021-5-28']
    # 训练自相关模型
    lr = evalute_autoregression(train, test, n_in=2, verbose=False)
    # 将数据集拆分为 X，y
    X_train, y_train = train.iloc[:, :2], train.iloc[:, -1]
    X_test, y_test = test.iloc[:, :2], test.iloc[:, -1]
    time_range = y_train.index
    
    # 自回归模型在训练集中的预测结果
    y_train_pred = lr.predict(X_train)
    # 自相关模型在训练集中的残差数据
    train_resid = y_train - y_train_pred
    # 残差数据画图
    plot_residual_error(time_range, train_resid)
    
    # 残差数据
    train_resid = pd.Series(train_resid)
    train_resid.index = time_range
    # 残差的自相关图
    acf_plotting(train_resid, residual=True)
    pacf_plotting(train_resid, residual=True)
    
    # 测试集的残差数据
    y_test_pred = lr.predict(X_test)
    time_range = y_test.index
    test_resid = y_test - y_test_pred
    test_resid = pd.Series(test_resid)
    test_resid.index = time_range
    # 使用滑动平均处理残差数据
    train_resid_win = series_to_supervised(train_resid, n_in=1, n_out=1, dropnan=True)
    x_train_resid = train_resid_win.iloc[:, 0]
    y_train_resid = train_resid_win.iloc[:, 1]
    
    test_resid_win = series_to_supervised(test_resid, n_in=1, n_out=1, dropnan=True)
    x_test_resid = test_resid_win.iloc[:, 0]
    y_test_resid = test_resid_win.iloc[:, 1]
    
    # 残差模型（滑动平均模型建立）
    lr_residual = evalute_autoregression(train_resid_win, test_resid_win, n_in=1, verbose=False)
    
    # 残差模型预测结果
    y_train_resid_pred = lr_residual.predict(x_train_resid.values.reshape(-1, 1))
    y_test_resid_pred = lr_residual.predict(x_test_resid.values.reshape(-1, 1))
    
    # 自相关+滑动平均
    y_train_pred = y_train_pred[1:] + y_train_resid_pred
    y_test_pred = y_test_pred[1:] + y_test_resid_pred
    
    print('模型在训练集中的 MSE 为: ', mean_squared_error(y_train.iloc[1:], y_train_pred))
    print('模型在测试集中的 MSE 为: ', mean_squared_error(y_test.iloc[1:], y_test_pred))
    
    
    